library(R2jags)
library(xlsx)
library(openxlsx)
library(spatstat)
library(dplyr)
library(readr)
require(R2wd)
library(LaplacesDemon) # for Geweke.Diagnostic()
library(ggplot2)

options(max.print=999999999)
memory.limit(10000);memory.size(max=T)

getparams<-function(jagssum,parameters=c('mu','sigma','diff','disc','beta'),questions=questions){
  for (n in parameters){
    v<-data.frame(jagssum[grepl(n,dimnames(jagssum)[[1]]),"mean"])
    s<-data.frame(jagssum[grepl(n,dimnames(jagssum)[[1]]),"sd"])
    for (w in c(parameters,'\\[','\\]')){
      row.names(v)<-gsub(w,"",row.names(v));row.names(s)<-gsub(w,"",row.names(s))}
    v <- cbind(rownames(v),v);s<-cbind(rownames(s),s)
    if (n=='beta'){
      v<-data.frame(v[,1:2]);s<-data.frame(s[,1:2]);colnames(v)<-c('n','mean.beta');colnames(s)<-c('n','sd.beta');pq<-left_join(v,s) 
    } else {
      if (n=="mu"|n=="sigma"){
        colnames(v)<-c('t',paste0("mean.",n));colnames(s)<-c('t',paste0("sd.",n));pq<-left_join(v,s)
      } else {
        if (n=="diff"){
          colnames(v)<-c('qn',paste0('mean.',n));colnames(s)<-c('qn',paste0('sd.',n))
          if (nrow(v)==info$nquest){
            v$qn<-as.character(v$qn);s$qn<-as.character(s$qn)
          } else {
            v$b<-substr(v$qn,(nchar(as.character(v$qn))+1)-1,nchar(as.character(v$qn)))
            s$b<-substr(s$qn,(nchar(as.character(s$qn))+1)-1,nchar(as.character(s$qn)))
            row.names(v)<-c(1:nrow(v));row.names(s)<-c(1:nrow(s))
            v$qn<-gsub(",1|,2|,3","",v$qn)
            s$qn<-gsub(",1|,2|,3","",s$qn)}
        } else {
          colnames(v)<-c('qn',paste0('mean.',n));colnames(s)<-c('qn',paste0('sd.',n))
          v$qn<-as.character(v$qn);s$qn<-as.character(s$qn)}
        pq<-left_join(v,s);pq<-left_join(pq,questions);pq$question<-as.character(pq$question)}}
    return(pq)}}

## JAGS code here
main_model="data{for (i in 1:len){opinionr[i,2]<-round(support[i]*samplesize[i]/100) 
		opinionr[i,1]<-round(neutral[i]*samplesize[i]/100)}}
model{for (i in 1:len){for(j in 1:2){opinionr[i,j]~dbin(p[i,j],samplesize[i])      
		p[i,j]~dbeta(alpha[i,j], beta)              
		alpha[i,j]<--beta*m[i,j]/(m[i,j]-1)         
		m[i,j]<-phi(x[i,j])             
	  x[i,j]<-(mu[t[i]]-diff[qn[i],j])/sqrt((disc[qn[i]])^2+(sigma[t[i]])^2)}}
beta~dunif(0,100)
mu[begin]~dnorm(0,1)
sigma_mu~dunif(0,10)
for (i in 1:nquest){disc[i]~dunif(0,6) 
	for (j in 1:2){diff0[i,j]~dunif(-4,4)}
	diff[i,1:2]<-sort(diff0[i,])}
for (i in begin:end){sigma[i]~dunif(0,10)}
for (i in begin:(end-1)){mu_raw[i]~dnorm(0,1)}
for (i in (begin+1):end){mu[i]~dnorm(mu[i-1]+sigma_mu*mu_raw[i-1],100)}}"
## Save to temp file
fileConn=file("main_model.temp")
writeLines(main_model,fileConn)
close(fileConn)

setwd(dirname(rstudioapi::getActiveDocumentContext()$path)) # set wd to source file location

outdir<-'irt_estimates/';dir.create(file.path(getwd(),outdir),showWarnings=F)
inpdir<-'raw_indicators/'

sheet<-read.table(paste0(inpdir,'EU_irt.txt'),sep=';',header=T)  # Read in data
sheet<-subset(sheet,as.numeric(gsub(' S1| S2','',sheet$t))>=1973)
sheet$t<-as.numeric(sheet$t)
sheet$asked<-NULL
sheet<-suppressMessages(left_join(sheet,sheet %>% group_by(question) %>% tally() %>% as.data.frame()))
names(sheet)[names(sheet)=='n']<-'asked'  
sheet<-subset(sheet,sheet$asked>1) 
sheet<-subset(sheet,sheet$question!='trustEU' & sheet$question!='trustECB'  & sheet$question!='trustEP' & sheet$question!='trustCOM' & sheet$question!='trustCOU' & sheet$question!='trustEUCO')

wb<-paste0(outdir,'EU_irt_estimates.xlsx')

rownames(sheet) <- 1:nrow(sheet);sheet$question<-factor(sheet$question) # Reset row numbers and update levels of q names after subsetting data
sheet$qn<-as.numeric(sheet$question);sheet<-sheet[order(sheet$qn,sheet$t),] # Sort by q number and time
questions<-unique(sheet[c("qn","question")]);questions$question<-as.character(questions$question);questions$qn<-as.character(questions$qn)
info<-list(len=nrow(sheet),nquest=length(unique(sheet$qn)),begin=min(sheet$t),end=max(sheet$t));data <-c(sheet,info)
  
jagsfit<-jags.parallel(data=data,parameters.to.save=c('mu','sigma',"diff","disc","beta"),
                       n.burnin=10000,n.iter=60000,n.thin=2,n.chains=3,
                       model.file="main_model.temp",jags.seed=123)

fit.mcmc<-as.mcmc(jagsfit);print(gelman.diag(fit.mcmc,autoburnin=F,transform=F)$mpsrf)

for (par in c('mu','sigma','diff','disc','beta')){p.func<-getparams(jagsfit$BUGSoutput$summary,par,questions); assign(par,p.func)};print(" ")
  
df<-merge(mu,sigma);df$t<-as.numeric(as.character(df$t));df<-df[order(df$t),] 
if ((nrow(diff))==info$nquest){diff<-diff} else {j=as.character(max(diff$b));diff<-subset(diff,b==j)}
for (r in 1:nrow(df)){
  df$Opinion[r]<-pnorm((df$mean.mu[r]-mean(diff$mean.diff))/(mean(disc$mean.disc)^2+df$mean.sigma[r]^2)^.5,0,1)*100
  df$Opinion_minus2sd[r]<-pnorm(((df$mean.mu[r]-2*df$sd.mu[r])-mean(diff$mean.diff))/(mean(disc$mean.disc)^2+df$mean.sigma[r]^2)^.5,0,1)*100
  df$Opinion_plus2sd[r]<-pnorm(((df$mean.mu[r]+2*df$sd.mu[r])-mean(diff$mean.diff))/(mean(disc$mean.disc)^2+df$mean.sigma[r]^2)^.5,0,1)*100
  df$Opinion_minussigma[r]<-df$Opinion[r]-(mean(disc$mean.disc)^2+df$mean.sigma[r]^2)^.5
  df$Opinion_plussigma[r]<-df$Opinion[r]+(mean(disc$mean.disc)^2+df$mean.sigma[r]^2)^.5}
if (!file.exists(wb)){workbook<-createWorkbook();addWorksheet(workbook,'jags_notrust');writeData(workbook,sheet='jags_notrust',x=df,rowNames=F)
  saveWorkbook(workbook,wb,overwrite=T)} else {workbook<-loadWorkbook(wb)
  if ('jags_notrust' %in% names(workbook)){removeWorksheet(workbook,'jags_notrust');addWorksheet(workbook,'jags_notrust')} else {addWorksheet(workbook,'jags_notrust')}
  writeData(workbook,sheet='jags_notrust',x=df,rowNames=F);saveWorkbook(workbook,wb,overwrite=T)}
